package com.caomeisan.muzi.kcface;

import android.text.TextUtils;
import android.util.Log;

import com.kuangchi.faiss.utils.FaissJni;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * @ClassName FaceInitUtils
 * @Description TODO
 * @Date 2023/11/2813:19
 * @Version 1.0
 */
public class FaceInitUtils {


    private static final String path = "/sdcard/SourceInfo/";
    private static final String TYPE_FLATIP = "FlatIP";
    private static final String TYPE_IVFFLATIP = "IVFFlatIP";
    private static final String TYPE_IVFPQ = "IVFPQ";
    private static final String TAG = "muzi102";
    public static final int FEATURE_SIZE = 1024;
    private FaissJni faissJni;
    private int nb = 0;// 向量集大小
    private int dim = 256;      // 向量维度
    private long faiss_index = 0;
    float[] mData = null;
    private List<FaceRegisterInfo> faceRegisterInfoList = new ArrayList<>();


    public FaceInitUtils() {
        faissJni = new FaissJni();
        initFile();
//        process(TYPE_IVFFLATIP);
        process(TYPE_FLATIP);
    }

    private void initFile() {
        File file = new File("/sdcard/faceSourceInfo/"); //存放体征的文件夹
        if (!file.exists()) {
            boolean b = file.mkdirs();
        }
        File path1 = new File(path); //存放 faces_%s.index features.dat 文件夹
        if (!path1.exists()) {
            path1.mkdirs();
        }

        File[] featureFiles = file.listFiles();
        if (featureFiles != null) {
            for (File featureFile : featureFiles) {
                try {
                    FileInputStream fis = new FileInputStream(featureFile);
                    byte[] feature = new byte[FEATURE_SIZE];
                    fis.read(feature);
                    faceRegisterInfoList.add(new FaceRegisterInfo(feature, featureFile.getName()));
                    fis.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            nb = featureFiles.length;
        }
    }

    private void process(String indexType) {

        int[] ids = new int[nb];
        float[] data = null;
        String index_file_path = path + "faces_%s.index";
        String features_file_path = path + "features.dat";

        data = load_features(nb, dim, features_file_path, ids);
        index_file_path = String.format(index_file_path, indexType);
        faiss_index = load_index(nb, dim, indexType, index_file_path, ids, data);
        mData = data;


    }


    //加载本地的特征
    private float[] load_features(int nb, int d, String features_file_path, int[] ids) {
        float[] data = null;
        File featuresFile = new File(features_file_path);
        if (featuresFile.exists()) {
            // 从文件中加载id和data
            data = new float[nb * d];
            readFeatures(nb, d, featuresFile, ids, data);
        } else {


            File newFile = new File(features_file_path);
            if (!newFile.exists()) {
                try {
                    newFile.createNewFile();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            float[] result = new float[nb * d];


            for (int i = 0; i < nb; i++) {

                float[] faceFloat = ByteUtils.byte2Float(faceRegisterInfoList.get(i).getFeatureData());//特征转float
                float[] faceFloatGy = ByteUtils.normalizeNew01(faceFloat); //归一化


                Log.e("muzi102","faceFloat = "+faceFloat.length);
                Log.e("muzi102","faceFloatGy = "+faceFloatGy.length);

                for (int j = 0; j < d; j++) {
                    result[i * d + j] = faceFloatGy[j];
                }


            }


            //data = ByteUtils.normalizeNew01(result);
            data = result;

            for (int i = 0; i < nb; i++) {
                ids[i] = i + 100;
            }
            // 保存id和data到文件

            writeFeatures(nb, d, featuresFile, ids, data);
        }
        return data;
    }

    private void readFeatures(int nb, int d, File featuresFile, int[] ids, float[] data) {
        try (DataInputStream inputStream = new DataInputStream(new BufferedInputStream(new FileInputStream(featuresFile)))) {

            for (int i = 0; i < nb; i++) {
                ids[i] = inputStream.readInt();
                for (int j = 0; j < d; j++) {
                    data[i * d + j] = inputStream.readFloat();
                }
            }
        } catch (IOException e) {
        }
    }

    private void writeFeatures(int nb, int d, File featuresFile, int[] ids, float[] data) {

        if (!featuresFile.exists()) {
            try {
                featuresFile.createNewFile();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        try (DataOutputStream outputStream = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(featuresFile)))) {
            for (int i = 0; i < nb; i++) {
                outputStream.writeInt(ids[i]);
                for (int j = 0; j < d; j++) {
                    outputStream.writeFloat(data[i * d + j]);
                }
            }
        } catch (IOException e) {
        }
    }

    private long load_index(int nb, int d, String type, String index_file_path, int[] ids, float[] data) {
        long t_index = -1;
        File indexFile = new File(index_file_path);
//        logger.info("indexFile:{}", indexFile.getAbsolutePath());
        if (indexFile.exists()) {
//            logger.info("read faiss index");
            t_index = faissJni.readIndexJNI(index_file_path);
        } else {
            File newFile = new File(index_file_path);
            if (!newFile.exists()) {
                try {
                    newFile.createNewFile();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if ("IVFFlatIP".equals(type)) {
                long nlist = 1024;  // 聚类中心的个数
                long nprobe = 16;   // 查找聚类中心的个数，默认为1个。
                t_index = faissJni.indexIVFFlatJNI(d, nlist, nprobe);
            } else if ("IVFPQ".equals(type)) {
                long nlist = 512;  // 聚类中心的数量
                long m = 64;        // 每个向量被划分为m个子向量
                long nbits = 8;     // 每个子向量使用nbits位进行量化
                t_index = faissJni.indexIVFPQJNI(d, nlist, m, nbits);
            } else {
                t_index = faissJni.indexFlatIPJNI(d);
            }
            // 在一些索引结构（比如IndexIVF及其变种）中，需要在索引之前对向量进行训练，以便索引能够更好地适应数据的分布。
            if ("IVFFlatIP".equals(type) || "IVFPQ".equals(type)) {
//                logger.info("train data");
                faissJni.trainJNI(t_index, nb, data);
            }
//            logger.info("add data");
            // 向索引中添加向量
            faissJni.addJNI(t_index, nb, data);
//            logger.info("write index");
            // 将索引保存到文件
            faissJni.writeIndexJNI(t_index, index_file_path);
        }
        return t_index;
    }

    //搜索
    public void search(byte[] features) {

        // 执行nq次查询并计算耗时和topK正确率
        long total_time = 0l;
        int query_vector_count = 1;
        float[] query_vector = ByteUtils.byte2Float(features);
//        Log.e(TAG, "特征float[] = " + Arrays.toString(query_vector));
        //归一化
        float[] query_vectorNew = ByteUtils.normalizeNew01(query_vector);
//        Log.e(TAG, Arrays.toString(query_vector));

        long start_time = System.currentTimeMillis();
        // 查找最相似的topK个向量

        String resultStr = faissJni.searchJNI(faiss_index, query_vector_count, query_vectorNew, 1);

        long end_time = System.currentTimeMillis();
        total_time = end_time - start_time;
        if (!TextUtils.isEmpty(resultStr)) {
            String[] resultArray = resultStr.split(",");
            String[] topOne = resultArray[0].split(":");
            long query_index = Long.parseLong(topOne[0]);
            Log.e(TAG, " 命中：query_index  = " + query_index + "" + " 命中：目标向量  = " + Float.parseFloat(topOne[1]) + " 命中：相似度 = " + ByteUtils.similarity(Float.parseFloat(topOne[1])));
        } else {
            Log.e(TAG, "未命中：目标向量{},匹配结果为空 ");
        }
        // 计算top1正确率
        Log.e(TAG, "总耗时: 毫秒 = " + total_time);
    }

    public List<FaceRegisterInfo> getFaceRegisterInfoList() {
        return faceRegisterInfoList;
    }

}
